package org.dxy.trigger;

import org.nd4j.linalg.factory.Nd4j;
import org.nd4j.linalg.factory.Nd4jBackend;
import org.nd4j.linalg.api.ndarray.INDArray;

/**
 * GPU测试类 - 验证CUDA配置是否正确
 */
public class GPUTest {
    public static void main(String[] args) {
        System.out.println("=== GPU Configuration Test ===");
        
        try {
            // 强制使用GPU后端
            System.setProperty("org.nd4j.linalg.factory", "org.nd4j.linalg.gpu.CudaNDArrayFactory");
            System.setProperty("org.nd4j.linalg.ops", "org.nd4j.linalg.gpu.ops.CudaOpExecutioner");
            
            // GPU内存管理
            System.setProperty("org.bytedeco.javacpp.maxbytes", "8G");
            System.setProperty("org.bytedeco.javacpp.maxphysicalbytes", "12G");
            
            // 获取后端信息
            Nd4jBackend backend = Nd4j.getBackend();
            System.out.println("ND4J Backend: " + backend.getClass().getName());
            
            // 检查GPU设备
            int deviceCount = Nd4j.getAffinityManager().getNumberOfDevices();
            System.out.println("Available GPU devices: " + deviceCount);
            
            if (deviceCount > 0) {
                System.out.println("GPU devices found:");
                for (int i = 0; i < deviceCount; i++) {
                    System.out.println("  Device " + i + " available");
                }
                
                // Get current thread device
                try {
                    Integer currentDevice = Nd4j.getAffinityManager().getDeviceForCurrentThread();
                    System.out.println("Current thread using device: " + currentDevice);
                } catch (Exception e) {
                    System.out.println("Device info: " + e.getMessage());
                }
                
                // GPU environment is ready
                System.out.println("GPU Environment ready for testing");
                
                // 测试GPU计算
                System.out.println("\n=== GPU Performance Test ===");
                testGPUPerformance();
                
            } else {
                System.out.println("No GPU devices found - falling back to CPU");
                testCPUFallback();
            }
            
        } catch (Exception e) {
            System.err.println("GPU test failed: " + e.getMessage());
            e.printStackTrace();
            testCPUFallback();
        }
    }
    
    private static void testGPUPerformance() {
        try {
            // 创建大矩阵进行GPU计算测试
            int size = 1000;
            System.out.println("Creating " + size + "x" + size + " matrices for GPU test...");
            
            long startTime = System.nanoTime();
            
            INDArray a = Nd4j.rand(size, size);
            INDArray b = Nd4j.rand(size, size);
            
//            System.out.println("Matrix A device: " + a.data().getDeviceId());
//            System.out.println("Matrix B device: " + b.data().getDeviceId());
            
            // 矩阵乘法
            INDArray result = a.mmul(b);
            
            long endTime = System.nanoTime();
            double duration = (endTime - startTime) / 1_000_000.0; // Convert to milliseconds
            
//            System.out.println("Result device: " + result.data().getDeviceId());
            System.out.println("Matrix multiplication completed in: " + String.format("%.2f", duration) + "ms");
            System.out.println("Result shape: " + java.util.Arrays.toString(result.shape()));
            System.out.println("First element: " + result.getDouble(0, 0));
            
            System.out.println("\n✅ GPU test completed successfully!");
            
        } catch (Exception e) {
            System.err.println("GPU performance test failed: " + e.getMessage());
            e.printStackTrace();
        }
    }
    
    private static void testCPUFallback() {
        try {
            System.out.println("\n=== CPU Fallback Test ===");
            
            // 强制使用CPU后端
            System.setProperty("org.nd4j.linalg.factory", "org.nd4j.linalg.cpu.nativecpu.CpuNDArrayFactory");
            System.setProperty("org.nd4j.linalg.ops", "org.nd4j.linalg.cpu.nativecpu.ops.NativeCpuOpExecutioner");
            
            // 重新初始化后端
            Nd4j.factory().setOrder('c');
            
            Nd4jBackend backend = Nd4j.getBackend();
            System.out.println("Fallback Backend: " + backend.getClass().getName());
            
            // 简单的CPU测试
            INDArray a = Nd4j.ones(100, 100);
            INDArray b = Nd4j.ones(100, 100);
            INDArray result = a.add(b);
            
            System.out.println("CPU test result: " + result.getDouble(0, 0));
            System.out.println("✅ CPU fallback working correctly!");
            
        } catch (Exception e) {
            System.err.println("CPU fallback test failed: " + e.getMessage());
            e.printStackTrace();
        }
    }
}